package com.ark.frigate.shell.utils;

import com.ark.frigate.shell.utils.dto.ConfigPo;
import com.ark.frigate.shell.utils.dto.ResultDto;
import com.jcraft.jsch.*;
import lombok.Builder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.*;
import java.nio.charset.StandardCharsets;

/**
 * 执行Shell脚本 或者 命令
 *
 * @author zengweilong
 * @date 5/31/21 2:27 PM
 */
@Slf4j
@Builder
public class JSchUtils {

    private ConfigPo configPo;

    private Session session;

    public boolean connect() {
        if (!checkParam()) {
            return false;
        }
        JSch jSch = new JSch();
        try {
            session = jSch.getSession(configPo.getUsername(), configPo.getHost(), configPo.getPort());
            if (StringUtils.isNoneBlank(configPo.getIdentity(), configPo.getPassword())) {
                jSch.addIdentity(configPo.getIdentity(), configPo.getPassword());
            } else if (StringUtils.isBlank(configPo.getIdentity())) {
                session.setPassword(configPo.getPassword());
            } else {
                jSch.addIdentity(configPo.getIdentity());
            }
            session.setConfig("PreferredAuthentications", "password");
            session.setConfig("StrictHostKeyChecking", "no");
            session.setTimeout(30000);
            session.connect();
            boolean result = session.isConnected();
            log.info("Connect result is = {}", result);
            return result;
        } catch (JSchException e) {
            log.error("JSch Exception Connect fail. configPo={} ", configPo, e);
        }
        return false;
    }

    /**
     * 脚本是同步执行的方式
     * 执行脚本命令
     *
     * @param command
     * @return
     */
    public ResultDto execCommand(String command) {
        ResultDto resultDto = new ResultDto();
        ChannelExec channel = null;

        StringBuffer result = new StringBuffer();
        BufferedReader reader = null;
        try {
            channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(command);
            channel.setInputStream(null);
            channel.setErrStream(System.err);
            channel.connect();
            InputStream in = channel.getInputStream();
            reader = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8));
            String res = "";
            while ((res = reader.readLine()) != null) {
                result.append(res + "\n");
                log.debug(res);
            }
            in.close();

            log.info("IP=[{} >>> command=[{}] >>> result=[{}]", configPo.getHost(), command, channel.getExitStatus());

            resultDto.setCode(channel.getExitStatus());
            resultDto.setResult(result.toString());
        } catch (Exception e) {
            log.error("Execute exec fail .", e);
        } finally {
            try {
                reader.close();
                closeChannel(channel);
            } catch (IOException e) {
                log.error(e.getMessage(), e);
            }
        }
        return resultDto;
    }


    public ResultDto executeShell(String command) {
        ResultDto resultDto = new ResultDto();
        ChannelShell channel = null;

        StringBuffer result = new StringBuffer();
        BufferedReader reader = null;
        InputStream in = null;
        OutputStream os = null;
        try {
            channel = (ChannelShell) session.openChannel("shell");
            channel.setPty(true);
            channel.connect();
            in = channel.getInputStream();
            os = channel.getOutputStream();
            os.write((command + "\r\n").getBytes());
            os.flush();
            reader = new BufferedReader(new InputStreamReader(in));
            String res = "";
            while ((res = reader.readLine()) != null) {
                result.append(res + "\n");
            }
            log.info("IP=[{} >>> command=[{}] >>> result=[{}]", configPo.getHost(), command, channel.getExitStatus());

            resultDto.setCode(channel.getExitStatus());
            resultDto.setResult(result.toString());
        } catch (Exception e) {
            log.error("Execute shell fail .", e);
        } finally {
            try {
                os.close();
                in.close();
                reader.close();
                closeChannel(channel);
            } catch (IOException e) {
                log.error(e.getMessage(), e);
            }
        }
        return resultDto;
    }

    /**
     * 将文件流 input 上传到远程路径 directory/sftpFileName
     *
     * @param directory
     * @param sftpFileName
     * @param input
     * @return
     */
    public ResultDto sftpUpload(String directory, String sftpFileName, InputStream input) {
        ResultDto resultDto = new ResultDto();
        ChannelSftp channel = null;
        String logFile = directory + sftpFileName;
        try {
            resultDto = execCommand("mkdir -p " + directory);
            log.info("create directory={}, result={}", directory, resultDto);
            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();
            channel.cd(directory);
            channel.put(input, directory + sftpFileName);
            resultDto.setCode(channel.getExitStatus() == -1 ? 0 : channel.getExitStatus());
            resultDto.setResult(channel.getExitStatus() == -1 ? "Upload Success:" + logFile : "Upload Error:" + logFile);
        } catch (Exception e) {
            resultDto.setCode(-2);
            resultDto.setResult("Execute sftp fail.");
            log.error("Execute sftp fail. {}", logFile, e);
        } finally {
            closeChannel(channel);
        }
        return resultDto;
    }

    /**
     * 关闭连接
     */
    private void closeChannel(Channel channel) {
        if (channel != null && channel.isConnected()) {
            channel.disconnect();
        }
    }

    /**
     * 关闭Session
     */
    public void closeSession() {
        if (session != null && session.isConnected()) {
            session.disconnect();
            session = null;
        }
    }

    private boolean checkParam() {

        if (StringUtils.isAnyBlank(configPo.getHost(), configPo.getUsername())) {
            log.warn(" connect host or username is null");
            return false;
        }


        if (null == configPo.getPort()) {
            log.warn(" connect post is null");
            return false;
        }

        if (StringUtils.isNoneBlank(configPo.getPassword(), configPo.getIdentity())) {
            log.warn(" connect password or identity is null");
            return false;
        }

        return true;
    }
}
